#include "refine/smooth_e.h"
#include "refine/varmap.h"
#include <limits>
#include <iostream>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth_e {

const params default_params = {2, 2, 1, 100, false};

typedef grey_mask::mask::elem_type grey_mask_value;
typedef unsigned uint;

const uint uchar_max = std::numeric_limits<unsigned char>::max();
const uint uchar_avg = uchar_max / 2;

/// describes set of pixels which are situated near pixel of interest and have similar colors
struct kin
{
public:
	size_point topleft;
	img::bitmap map;
public:
	template <typename Raster>
	void create(Raster const & im, size_point const & c, uint whs);
	template <typename Matte>
	grey_mask_value classify(Matte const & matte, size_point const & c) const;
};

typedef std::vector<kin> kins;

inline uint dist(int a, int b) 
	{ return (a - b) * (a - b); }

inline uint dist(size_point const & a, size_point const & b) 
	{ return dist(a.x, b.x) + dist(a.y, b.y); }

template <typename Col>
inline uint cdist(Col const & a, Col const & b) 
	{ return dist(a.b, b.b) + dist(a.g, b.g) + dist(a.r, b.r); }

template <typename Raster>
void kin::create(Raster const & im, size_point const & c, uint whs)
{
	topleft.x = c.x < whs ? 0 : c.x - whs;
	topleft.y = c.y < whs ? 0 : c.y - whs;

	size_extent ext;
	ext.sx = c.x + whs + 1 >= im.width()  ? im.width()  - topleft.x : c.x + whs + 1 - topleft.x;
	ext.sy = c.y + whs + 1 >= im.height() ? im.height() - topleft.y : c.y + whs + 1 - topleft.y;
	map.resize(ext);
	map.fill(false);

	typedef std::pair<uint, size_point> dist_pos; //(color distance, position in map)
	std::vector<dist_pos> neighbours;
	neighbours.reserve(ext.sx * ext.sy);
	size_point p;
	typename Raster::elem_type const & v0 = im[c];
	for (p.y = 0; p.y < ext.sy; ++p.y)
		for (p.x = 0; p.x < ext.sx; ++p.x) {
			if (whs > 2 && dist(topleft + p, c) > whs * whs)
				continue; //use circle window
			typename Raster::elem_type const & v = im[topleft + p];
			neighbours.push_back(dist_pos(cdist(v0, v), p));
		}

	uint count = std::min(neighbours.size(), std::max(7u, neighbours.size() / 3));
	std::partial_sort(neighbours.begin(), neighbours.begin() + count, neighbours.end());
	//mark selected pixels
	for (uint i = 0; i < count; ++i)
		map[neighbours[i].second] = true;
}

template <typename Matte>
grey_mask_value kin::classify(Matte const & matte, size_point const & c) const
{
	//compute mean alpha among kin
	uint sum = 0;
	uint count = 0;
	size_point p;
	for (p.y = 0; p.y < map.height(); ++p.y)
		for (p.x = 0; p.x < map.width(); ++p.x) {
			if (!map[p])
				continue;
			sum += matte[topleft + p];
			++count;
		}
	return (sum + count/2) / count;
}

template <typename Raster, typename Matte>
static void perform_core(
	const Raster & src, 
	Matte & matte, 
	const img::bitmap & strict_fore, 
	const img::bitmap & strict_back,
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Volosok(R) matting..." << endl;

	size_extent dim = src.dim();
	assert (dim == matte.dim());
	assert (dim == modifiable.dim());
	int_rect frame(0, 0, dim.sx, dim.sy);

	varmap vmap(strict_fore, strict_back);
	kins ks(vmap.count());
	size_point c;
	for (c.y = 0; c.y < src.height(); ++c.y)
		for (c.x = 0; c.x < src.width(); ++c.x)
			if (vmap.variable(c))
				ks[vmap[c]].create(src, c, p.whs);

	grey_mask::mask next_matte = matte;

	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		img::combine(modifiable, strict_fore, img::bit_and_not());
		img::combine(modifiable, strict_back, img::bit_and_not());

		uint check_count = 0, change_count = 0;

		for (c.y = 0; c.y < src.height(); ++c.y)
			for (c.x = 0; c.x < src.width(); ++c.x)
				if (modifiable[c]) {
					++check_count;
					kin const & k = ks[vmap[c]];
					grey_mask_value g = k.classify(matte, c);
					if (matte[c] == g)
						modifiable[c] = false;
					else {
						next_matte[c] = g;
						++change_count;
					}
				}

		if (!p.silent)
			clog << ": checked = " << check_count << ", alpha changes = " << change_count << endl;

		if (change_count == 0)
			break;

		matte = next_matte;

		img::dilate8(modifiable, p.whs);
	}
}

template <typename Raster, typename Matte>
static void perform_from_bit_mask(
	const Raster & src, 
	const bit_mask::mask & m, 
	Matte & matte, 
	const params & p)
{
	size_extent dim = src.dim();
	assert (dim == m.dim());
	assert (dim == matte.dim());

	for (typename Matte::iterator i = matte.begin(); i != matte.end(); ++i)
		*i = m[i] ? uchar_max : 0;

	img::bitmap strict_fore, strict_back;
	img::erode4(strict_fore, m, p.fore_depth);
	img::dilate4(strict_back, m, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable = m; 
	{
		uint w = 1 + p.whs;
		img::dilate4(modifiable, w);
		img::bitmap tmp = m;
		tmp.inverse();
		img::dilate4(tmp, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

template <typename Raster, typename Matte>
static void perform_from_grey_mask(
	const Raster & src, 
	Matte & matte, 
	const params & p)
{
	size_extent dim = matte.dim();

	img::bitmap not_min(dim), not_max(dim), strict_fore(dim), strict_back(dim);
	for (typename Matte::/*const_*/iterator i = matte.begin(); i != matte.end(); ++i) {
		not_min[i] = *i != 0;
		not_max[i] = *i != uchar_max;
		strict_fore[i] = *i > uchar_avg;
	}

	strict_back = strict_fore;
	img::erode4(strict_fore, p.fore_depth);
	img::dilate4(strict_back, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs;
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

// exported functions

void perform(
	const byte_image & src, 
	const bit_mask::mask & m, 
	grey_mask::mask & matte, 
	const params & p)
{
	matte.resize(src.dim());
	perform_from_bit_mask(src, m, matte, p);
}

void perform(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const params & p)
{
	matte.resize(src.dim());
	perform_from_grey_mask(src, matte, p);
}

void perform(
	byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, img::projection(im, &byte_color::a), p);
}

void perform(
	argb_byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, img::projection(im, &argb_byte_color::a), p);
}

} //namespace smooth_e

} //namespace refine
